import pandas as pd
import numpy as np
from castle.metrics import MetricsDAG
from castle.algorithms import TTPM

if __name__ == '__main__':
    file_index = 7
    path = './datasets_phase2/'
    alarm_data = pd.read_csv(path + str(file_index) + '/Alarm.csv', encoding='utf')
    alarm_data = alarm_data.iloc[:, 0:3]
    alarm_data.columns = ['event', 'node', 'timestamp']
    alarm_data = alarm_data.reindex(columns=['event', 'timestamp', 'node'])
    alarm_data = alarm_data.iloc[:10000, :]
    # topology_matrix = np.load(path + str(file_index) + '/Topology.npy')
    topology_matrix = np.zeros(shape=(alarm_data["node"].max() + 1, alarm_data["node"].max() + 1))

    ttpm = TTPM(topology_matrix, max_hop=2, max_iter=100, delta=0.01, penalty='BIC', epsilon=0.6)

    ttpm.learn(alarm_data)
    est_causal_matrix = np.array(ttpm.causal_matrix)
    np.fill_diagonal(est_causal_matrix, 0)
    np.save('./est/' + str(file_index) + '.npy', est_causal_matrix)
